// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/functions/registerer.h>
#include <pollux/functions/lib/re2_functions.h>
#include <pollux/functions/prestosql/regexp_replace.h>
#include <pollux/functions/prestosql/split_part.h>
#include <pollux/functions/prestosql/split_to_map.h>
#include <pollux/functions/prestosql/split_to_multi_map.h>
#include <pollux/functions/prestosql/string_functions.h>
#include <pollux/functions/prestosql/word_stem.h>

namespace kumo::pollux::functions {
    namespace {
        std::shared_ptr<exec::VectorFunction> makeRegexExtract(
            const std::string &name,
            const std::vector<exec::VectorFunctionArg> &inputArgs,
            const core::QueryConfig &config) {
            return makeRe2Extract(name, inputArgs, config, /*emptyNoMatch=*/false);
        }

        void registerSimpleFunctions(const std::string &prefix) {
            using namespace stringImpl;

            // Register string functions.
            register_function<ChrFunction, Varchar, int64_t>({prefix + "chr"});
            register_function<CodePointFunction, int32_t, Varchar>({prefix + "codepoint"});
            register_function<HammingDistanceFunction, int64_t, Varchar, Varchar>(
                {prefix + "hamming_distance"});
            register_function<LevenshteinDistanceFunction, int64_t, Varchar, Varchar>(
                {prefix + "levenshtein_distance"});
            register_function<LengthFunction, int64_t, Varchar>({prefix + "length"});
            register_function<XxHash64StringFunction, int64_t, Varchar>(
                {prefix + "xxhash64_internal"});

            // Length for varbinary have different semantics.
            register_function<LengthVarbinaryFunction, int64_t, Varbinary>(
                {prefix + "length"});

            register_function<StartsWithFunction, bool, Varchar, Varchar>(
                {prefix + "starts_with"});
            register_function<EndsWithFunction, bool, Varchar, Varchar>(
                {prefix + "ends_with"});
            register_function<EndsWithFunction, bool, Varchar, UnknownValue>(
                {prefix + "ends_with"});

            register_function<TrailFunction, Varchar, Varchar, int32_t>(
                {prefix + "trail"});

            register_function<SubstrFunction, Varchar, Varchar, int64_t>(
                {prefix + "substr"});
            register_function<SubstrFunction, Varchar, Varchar, int64_t, int64_t>(
                {prefix + "substr"});

            // TODO Presto doesn't allow INTEGER types for 2nd and 3rd arguments. Remove
            // these signatures.
            register_function<SubstrFunction, Varchar, Varchar, int32_t>(
                {prefix + "substr"});
            register_function<SubstrFunction, Varchar, Varchar, int32_t, int32_t>(
                {prefix + "substr"});

            register_function<SubstrVarbinaryFunction, Varbinary, Varbinary, int64_t>(
                {prefix + "substr"});
            register_function<
                SubstrVarbinaryFunction,
                Varbinary,
                Varbinary,
                int64_t,
                int64_t>({prefix + "substr"});

            register_function<SplitPart, Varchar, Varchar, Varchar, int64_t>(
                {prefix + "split_part"});

            register_function<TrimFunction, Varchar, Varchar>({prefix + "trim"});
            register_function<TrimFunction, Varchar, Varchar, Varchar>({prefix + "trim"});
            register_function<LTrimFunction, Varchar, Varchar>({prefix + "ltrim"});
            register_function<LTrimFunction, Varchar, Varchar, Varchar>(
                {prefix + "ltrim"});
            register_function<RTrimFunction, Varchar, Varchar>({prefix + "rtrim"});
            register_function<RTrimFunction, Varchar, Varchar, Varchar>(
                {prefix + "rtrim"});

            register_function<LPadFunction, Varchar, Varchar, int64_t, Varchar>(
                {prefix + "lpad"});
            register_function<RPadFunction, Varchar, Varchar, int64_t, Varchar>(
                {prefix + "rpad"});

            exec::registerStatefulVectorFunction(
                prefix + "like", likeSignatures(), makeLike);

            register_function<Re2RegexpReplacePresto, Varchar, Varchar, Varchar>(
                {prefix + "regexp_replace"});
            register_function<Re2RegexpReplacePresto, Varchar, Varchar, Varchar, Varchar>(
                {prefix + "regexp_replace"});
            exec::registerStatefulVectorFunction(
                prefix + "regexp_replace",
                regexpReplaceWithLambdaSignatures(),
                makeRegexpReplaceWithLambda,
                exec::VectorFunctionMetadataBuilder().defaultNullBehavior(false).build());

            register_function<Re2RegexpSplit, Array<Varchar>, Varchar, Varchar>(
                {prefix + "regexp_split"});
        }

        void registerSplitToMultiMap(const std::string &prefix) {
            register_function<
                SplitToMultiMapFunction,
                Map<Varchar, Array<Varchar> >,
                Varchar,
                Varchar,
                Varchar>({prefix + "split_to_multimap"});
            register_function<
                SplitToMultiMapFunction,
                Map<Varchar, Array<Varchar> >,
                Varchar,
                UnknownValue,
                Varchar>({prefix + "split_to_multimap"});
            register_function<
                SplitToMultiMapFunction,
                Map<Varchar, Array<Varchar> >,
                Varchar,
                Varchar,
                UnknownValue>({prefix + "split_to_multimap"});
            register_function<
                SplitToMultiMapFunction,
                Map<Varchar, Array<Varchar> >,
                Varchar,
                UnknownValue,
                UnknownValue>({prefix + "split_to_multimap"});
        }

        void registerSplitToMap(const std::string &prefix) {
            register_function<
                SplitToMapFunction,
                Map<Varchar, Varchar>,
                Varchar,
                Varchar,
                Varchar>({prefix + "split_to_map"});

            register_function<
                SplitToMapFunction,
                Map<Varchar, Array<Varchar> >,
                Varchar,
                UnknownValue,
                Varchar>({prefix + "split_to_map"});
            register_function<
                SplitToMapFunction,
                Map<Varchar, Array<Varchar> >,
                Varchar,
                Varchar,
                UnknownValue>({prefix + "split_to_map"});
            register_function<
                SplitToMapFunction,
                Map<Varchar, Array<Varchar> >,
                Varchar,
                UnknownValue,
                UnknownValue>({prefix + "split_to_map"});

            exec::registerVectorFunction(
                prefix + "split_to_map",
                {
                    exec::FunctionSignatureBuilder()
                    .returnType("map(varchar,varchar)")
                    .argumentType("varchar")
                    .argumentType("varchar")
                    .argumentType("varchar")
                    .argumentType("function(varchar,varchar,varchar,varchar)")
                    .build(),
                },
                std::make_unique<exec::ApplyNeverCalled>());
            register_function<
                SplitToMapFunction,
                Map<Varchar, Varchar>,
                Varchar,
                Varchar,
                Varchar,
                bool>({"$internal$split_to_map"});
            exec::registerExpressionRewrite([prefix](const auto &expr) {
                return rewriteSplitToMapCall(prefix, expr);
            });
        }
    } // namespace

    void registerStringFunctions(const std::string &prefix) {
        registerSimpleFunctions(prefix);

        POLLUX_REGISTER_VECTOR_FUNCTION(udf_lower, prefix + "lower");
        POLLUX_REGISTER_VECTOR_FUNCTION(udf_upper, prefix + "upper");
        POLLUX_REGISTER_VECTOR_FUNCTION(udf_split, prefix + "split");

        registerSplitToMap(prefix);
        registerSplitToMultiMap(prefix);

        POLLUX_REGISTER_VECTOR_FUNCTION(udf_concat, prefix + "concat");
        POLLUX_REGISTER_VECTOR_FUNCTION(udf_replaceFirst, prefix + "replace_first");
        POLLUX_REGISTER_VECTOR_FUNCTION(udf_replace, prefix + "replace");
        POLLUX_REGISTER_VECTOR_FUNCTION(udf_reverse, prefix + "reverse");
        POLLUX_REGISTER_VECTOR_FUNCTION(udf_to_utf8, prefix + "to_utf8");
        POLLUX_REGISTER_VECTOR_FUNCTION(udf_from_utf8, prefix + "from_utf8");

        // Regex functions
        exec::registerStatefulVectorFunction(
            prefix + "regexp_extract", re2ExtractSignatures(), makeRegexExtract);
        exec::registerStatefulVectorFunction(
            prefix + "regexp_extract_all",
            re2ExtractAllSignatures(),
            makeRe2ExtractAll);
        exec::registerStatefulVectorFunction(
            prefix + "regexp_like", re2SearchSignatures(), makeRe2Search);

        register_function<StrLPosFunction, int64_t, Varchar, Varchar>(
            {prefix + "strpos"});
        register_function<StrLPosFunction, int64_t, Varchar, Varchar, int64_t>(
            {prefix + "strpos"});
        register_function<StrRPosFunction, int64_t, Varchar, Varchar>(
            {prefix + "strrpos"});
        register_function<StrRPosFunction, int64_t, Varchar, Varchar, int64_t>(
            {prefix + "strrpos"});

        register_function<NormalizeFunction, Varchar, Varchar>({prefix + "normalize"});
        register_function<NormalizeFunction, Varchar, Varchar, Varchar>(
            {prefix + "normalize"});

        // word_stem function
        register_function<WordStemFunction, Varchar, Varchar>({prefix + "word_stem"});
        register_function<WordStemFunction, Varchar, Varchar, Varchar>(
            {prefix + "word_stem"});
    }
} // namespace kumo::pollux::functions
